qwen4_exp: load the block-FP8 dense projections (fixes modelopt MIXED_PRECISION checkpoints) - #392

Ran this on 2 x RTX 6000 Ada (sm_89, 48 GB each, PCIe 4 x16, no NVLink) with the lovedheart NVFP4-FP8 (modelopt MIXED_PRECISION) checkpoint, offload backend, KV pool 262,144 tokens, pinned PLE, 8 running requests, one card (TP=1) unless stated. Tree: main af71ba4 + TP (#385) + load-time per-tensor FP8 (#389) + this PR cherry-picked. Same harness as my other PR feedback: single-stream = median of three 256-token generations, aggregate = eight concurrent 256-token requests, TTFT on a ~1k-token prompt, residency = expert slots / 24,576, probe = the same 8-question greedy smoke test (thinking off).

So the win reproduces on Ada at TP=1, and it is the residency effect you describe: +3.9 points of experts on the card. In isolation the Triton block-FP8 GEMV is slower than bf16 cuBLAS on this card (3.78 vs 3.20 ms per 48-layer decode step at M=1, the micro-benchmark in #389), but the freed 2.5 GiB buys more than that costs. Per-tensor _scaled_mm edges it out by 2-3% at the same residency because that GEMM is about 2x faster than the block kernel on sm_89.

At TP=2 (both cards, 16 running, vision tower loaded) the block-FP8 linears have no parallel variant, so I carry this PR on my deploy branch with a precedence rule: FREETOKEN_FP8_DENSE=1 dequantizes and re-quantizes per tensor (TP-capable); otherwise native block-FP8 at TP=1 (this PR); otherwise bf16 under TP>1. Lovedheart at TP=2 that way: 99.8 tok/s single-stream, 336.7 at 8 concurrent, 95.8% residency, TTFT 0.81 s, the same as the RadixArk NVFP4 build with the same flag (99.2-99.6 / 331.9). Probe 6/8 at TP=2. The 6/8 at TP=2 is the prime-count question (primes in 10..30, answered 5): the RadixArk build also misses it at TP=2 in bf16 and in FP8, the bf16-dequant lovedheart run at TP=2 scored 7/8, and every TP=1 run scored 7/8. It is a borderline question that flips with the numeric path (TP=2 reduction order, FP8), and this is a one-question smoke test, not a perplexity measurement. Greedy outputs against the bf16-dequant run: the 1k-token prompt is identical for both FP8 routes at TP=1; the two short prompts diverge after 13-79 words, against a run-to-run floor of ~84 words / identical / identical on this box.

Two notes on the diff from that merge: (1) routing the bf16 fallback of attention.py through make_col_merged / make_replicated drops the TP-aware classes (local_output_sizes, row-parallel o_proj) that #385 needs; in my merge the factories are used only for the fp8_block branch. (2) _dense_is_block_fp8(model_path) in weight.py and parse_config read the declaration independently; when a rank has to downgrade (TP>1) the two must agree, so I resolve both from one helper (dense_quant_mode) after the TP info is set. Branch: gdevenyi/FreeToken deploy/chatdnp (5d97d3a).

Thanks for taking both points. The dense_quant_mode consolidation matches what I carry on the deploy branch, and resolving it through one helper with try_get_tp_info is the right call — a rank has to build and load the same mode or the buffers do not match.

One thing to flag on point 1, not a problem with this PR as it stands but a hazard when it meets #385.

The new bf16 branch in models/qwen4_exp/attention.py hardcodes:

self.o_proj = LinearReplicated(self.qo_attn_dim, config.hidden_size, has_bias=False)

That is exactly what main does today, so this PR is correct on its own. But it is also the precise path a rank takes under TP>1 — the one your commit message calls out — and there LinearReplicated is wrong in three ways at once. qkv_proj is column-parallel, so each rank's attention output is its local head slice [T, local_qo_dim], not the full qo_attn_dim; o_proj therefore has to take the sharded input dim and all-reduce the partial sums afterwards. LinearReplicated keeps the full [hidden, qo_attn_dim] weight, expects the unsharded input, and performs no reduction. The loader already assumes the row-parallel layout: _shard puts o_proj on dim 1.

On my deploy branch the same else branch reads:

self.qkv_proj = LinearColParallelMerged(
    config.hidden_size, qkv_sizes, has_bias=False, local_output_sizes=self._qkv_split
)
self.o_proj = LinearOProj(self.qo_attn_dim, config.hidden_size, has_bias=False)

LinearOProj divides the input size by tp_info.size and all-reduces after the GEMM, and it degenerates to LinearReplicated's behaviour at TP=1 (div_even(x, 1) == x, the all-reduce is skipped when tp_size == 1), so it is safe for main as it stands. The one caveat is that it calls get_tp_info() in __init__, which raises when no engine has set TP info — the same reason you reached for try_get_tp_info in dense_quant_mode. So it is not a drop-in for tooling paths that build a model with no engine.

I am not asking you to change anything here; #385 is mine and the merge is my problem. Flagging it so whichever of the two lands second does not silently keep the LinearReplicated line, because the failure is quiet: the shapes only disagree once local_output_sizes shortens the qkv split, and a missing all-reduce gives each rank a partial sum that still decodes to fluent-looking text.

Also worth knowing, since your commit message mentions the two cannot be tested together on your box: the combination is running here. deploy/chatdnp carries #385 + #386 + #389 + this PR + #354, TP=2 on 2 x RTX 6000 Ada, and both attention branches are exercised — fp8_dynamic in production and the bf16 fallback whenever FREETOKEN_FP8_DENSE=0. If a TP=2 run of a specific case would help, say which and I will post it.

🤖 Generated with Claude Code

https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt

Thanks, that's a fair flag and cheaper to fix now than to remember at merge time. Swapped in 2b2a96e: the bf16 branch now builds o_proj as LinearOProj, same as your deploy branch. I read the class before doing it and agree it's a no-op at TP=1: div_even(x, 1) == x and the all-reduce is skipped, so the weight shape and the GEMM are identical to LinearReplicated. The comment above that branch also promised a row-parallel o_proj the code didn't build, so it now says what's there.

On the get_tp_info() caveat: it doesn't add a constraint here, because the same else branch already calls it two lines up through LinearColParallelMerged. That path was engine-only before and still is. Config parsing is the no-engine path that mattered, and that stays on try_get_tp_info. I've left local_output_sizes to #385, since main's LinearColParallelMerged doesn't take it, and likewise the forward's reshape(-1, self.qo_attn_dim) on the gate, which is the other full-width assumption in this file and belongs with whoever owns the local head count.

Verified at TP=1 on the Blackwell box: the four qwen4_exp test files that build this layer (test_config, test_weight, test_skeleton, test_qsa_backend) are 49/49 at 72773b0 and 49/49 with the change, including test_qsa_layer_matches_hf_dense, which runs the whole layer on the card against the HF dense reference. A direct check on CPU and CUDA confirms the built o_proj is a LinearOProj with weight [hidden, qo_attn_dim] and a forward bit-identical to LinearReplicated carrying the same weight. One correction to my last commit message while I'm here: the 47 failures I put down to VRAM contention were, for these four files at least, a missing ninja on PATH in my throwaway venv. With that fixed they pass in full.

I'll take you up on the TP=2 run. The case I can't exercise is exactly this one, the bf16 fallback under TP, so: deploy/chatdnp with this commit, FREETOKEN_FP8_DENSE=0, TP=2, your usual single-stream and the 8-question probe, with your fp8_dynamic production numbers as the reference. If the probe is still 7/8 and the rate lands where the fallback did before, that's the missing all-reduce ruled out rather than merely fluent.